package com.nulldev.util.web.dns.server.message;

import java.nio.ByteBuffer;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import com.nulldev.util.web.dns.server.util.ByteBufferUtils;

public class DNSQueryBuilder {
	private static Random random;

    static {
        try {
            random = SecureRandom.getInstanceStrong();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            random = new Random();
        }
    }

    private ByteBuffer messageData;
    private DNSMessageOverlay overlay;
    private List<DNSQuestion> questions = new ArrayList<DNSQuestion>();

    /**
     * Creates a new DnsQueryBuilder, initializing the id field to a random value.
     */
    public DNSQueryBuilder() {
        messageData = ByteBuffer.allocate(DNSMessageOverlay.MAX_UDP_DNS_LENGTH);
        overlay = new DNSMessageOverlay(messageData);
        overlay.setIdentifier((short) random.nextInt());
    }

    /**
     * Set the identifier field of a DNS query.
     * @param id
     * @return this
     */
    public DNSQueryBuilder setIdentifier(short id) {
        overlay.setIdentifier(id);
        return this;
    }

    /**
     * Set OP code of DNS query.
     * @param code
     * @return this
     */
    public DNSQueryBuilder setOperationCode(DNSOpCode code) {
        overlay.setOperationCode(code);
        return this;
    }

    /**
     * Indicates whether resolver should use recursion (i.e. travers the DNS hierarchy) when resolving questions.
     * @param desired
     * @return this
     */
    public DNSQueryBuilder setRecursionDesired(boolean desired) {
        overlay.setRecursionDesired(desired);
        return this;
    }

    /**
     * Adds a {@link org.code_revue.dns.message.DnsQuestion} to the query. This can be called multiple times when
     * building a query.
     * @param question
     * @return this
     */
    public DNSQueryBuilder addQuestion(DNSQuestion question) {
        questions.add(question);
        return this;
    }

    /**
     * Creates a {@link java.nio.ByteBuffer} with the binary representation of this DNS query.
     * @return Binary data for this query
     */
    public ByteBuffer build() {
        overlay.setQuery();
        overlay.setTruncated(false);
        overlay.setQuestionCount((short) questions.size());

        messageData.position(DNSMessageOverlay.HEADER_LENGTH);
        for (DNSQuestion question: questions) {
            ByteBufferUtils.writeDomainName(question.getQuestionName(), messageData);
            messageData.putShort((short) question.getQuestionType().getNumericCode());
            messageData.putShort((short) 1);
        }

        messageData.limit(messageData.position());
        messageData.position(0);
        return messageData.slice();
    }
}
